#include "record_k4a_image.h"
#include "compress.hpp"
#include <condition_variable>
#include <string>
#include <fstream>
#include <sstream>
#include <opencv2/opencv.hpp>

static std::condition_variable s_cv;
static std::mutex s_mtx;
static bool s_ready = false;
static BytesBuffer* s_preBuf = nullptr;

void RecordImage::compress(const char *in_buf, uint32_t in_buf_size, BytesBuffer* out)
{
    if (out->writableBytes() < in_buf_size)
        return;

    std::vector<char> in_vec;
    std::vector<char> out_vec;
    in_vec.assign(in_buf, in_buf + in_buf_size);

    gzip_compress(in_vec, out_vec);
    uint32_t len = out_vec.size();
    out->append((char*)&len, 4);
    out->append(&out_vec[0], len);
}

void RecordImage::compress(BytesBuffer* in, BytesBuffer* out)
{
    compress(in->peek(), in->readableBytes(), out);
}

void RecordImage::decompress(std::vector<char>& in_vec, BytesBuffer* out)
{
    std::vector<char> out_vec;
    gzip_decompress(in_vec, out_vec);
    out->append((char*)&out_vec[0], out_vec.size());
}

void RecordImage::decompress(BytesBuffer* buffer, BytesBuffer* out)
{
    uint32_t len = 0;
    if (buffer->readableBytes() < 4)
        return;
    buffer->retrieve((char*)&len, 4);
    if (len > buffer->readableBytes())
        return;
    std::vector<char> in_vec;
    in_vec.assign(buffer->peek(), buffer->peek() + len);
    buffer->retrieve(len);

    decompress(in_vec, out);
}

void RecordImage::showDebugImage(BytesBuffer* p, int imageType)
{
    int height_pixels;
    int width_pixels;
    int stride_bytes;
    p->retrieve((char*)&height_pixels, 4);
    p->retrieve((char*)&width_pixels, 4);
    p->retrieve((char*)&stride_bytes, 4);
    if (imageType == kImageType_rgb)
    {
        cv::Mat cv_rgbImage_no_alpha;
        cv::Mat cv_rgbImage_with_alpha = cv::Mat(height_pixels, width_pixels, CV_8UC4, (void*)(p->peek() + 4 * 3));
        cv::cvtColor(cv_rgbImage_with_alpha, cv_rgbImage_no_alpha, cv::COLOR_BGRA2BGR);
        cv::imshow("rgb", cv_rgbImage_no_alpha);
    }
    else if (imageType == kImageType_depth)
    {
        cv::Mat cv_depth_8U;
        cv::Mat img(height_pixels, width_pixels, CV_16U, (void*)(p->peek() + 4 * 3), stride_bytes);
        img.convertTo(cv_depth_8U, CV_8U, 1);
        cv::imshow("depth", cv_depth_8U);
    }
    else if (imageType == kImageType_ir)
    {
        cv::Mat cv_irImage_8U;
        cv::Mat img(height_pixels, width_pixels, CV_16U, (void*)(p->peek() + 4 * 3), stride_bytes);
        img.convertTo(cv_irImage_8U, CV_8U, 1);
        cv::imshow("ir", cv_irImage_8U);
    }
    
    cv::waitKey(10);
}

void write_depth_png(BytesBuffer* p)
{
    int height_pixels;
    int width_pixels;
    int stride_bytes;
    p->retrieve((char*)&height_pixels, 4);
    p->retrieve((char*)&width_pixels, 4);
    p->retrieve((char*)&stride_bytes, 4);
    cv::Mat cv_depth_8U;
    cv::Mat depthMat(height_pixels, width_pixels, CV_16U, (void*)(p->peek() + 4 * 3), stride_bytes);
    static int s_index = 1;
    std::string dst = std::to_string(s_index++) + ".png";
    cv::imwrite(dst, depthMat, { cv::IMWRITE_PNG_COMPRESSION , 0 });
}

void RecordImage::compressAndWriteImageData(k4a::image& img)
{
    if (_curBuf == nullptr)
    {
        _curBuf = &_buffer1;
    }
    int dataLen = 4 * 3 + img.get_size();
    if (_curBuf->writableBytes() < dataLen)
    {
        std::unique_lock<std::mutex> lck(s_mtx);
        if (_curBuf->readableBytes() > 0) // exist data
        { 
            s_preBuf = _curBuf;
            if (_curBuf == &_buffer1)
                _curBuf = &_buffer2;
            else
                _curBuf = &_buffer1;
            s_ready = true;
            lck.unlock();
            s_cv.notify_one();
            std::this_thread::yield();
        }
    }

    int height_pixels = img.get_height_pixels();
    int width_pixels = img.get_width_pixels();
    int stride_bytes = img.get_stride_bytes();

    std::vector<char> in_vec;
    BytesBuffer tmp_buf(dataLen);
    tmp_buf.append((char*)&height_pixels, 4);
    tmp_buf.append((char*)&width_pixels, 4);
    tmp_buf.append((char*)&stride_bytes, 4);
    tmp_buf.append((char*)img.get_buffer(), img.get_size());
    compress(&tmp_buf, _curBuf);
#if 0
    //////////////////
    BytesBuffer out;
    std::vector<char> vec;
    decompress(_curBuf, &out);

    showDebugImage(&out, kImageType_depth);
    /////////////////////////
#endif
}

void RecordImage::writeToFile()
{
    while (true)
    {
        std::unique_lock<std::mutex> lck(s_mtx);
        s_cv.wait(lck);
        if (!s_ready)
            continue;

        std::ofstream file;
        if (file.bad())
        {
            std::cerr << "cannot open file" << std::endl;
            break;
        }
        file.open(_fileName, std::ios::out | std::ios::binary | std::ios::app);
        file.write(s_preBuf->peek(), s_preBuf->readableBytes());
        file.close();
        s_preBuf->retrieveAll();
        s_ready = false;
    }
}
